from django.db import migrations, models

def update_industry_to_new_material(apps, schema_editor):
    """更新所有表的产业字段为新材料"""
    ProvincePatent = apps.get_model('trend', 'ProvincePatent')
    Keyword = apps.get_model('trend', 'Keyword')
    TopOrg = apps.get_model('trend', 'TopOrg')
    WordTrend = apps.get_model('trend', 'WordTrend')
    
    ProvincePatent.objects.all().update(industry='新材料')
    Keyword.objects.all().update(industry='新材料')
    TopOrg.objects.all().update(industry='新材料')
    WordTrend.objects.all().update(industry='新材料')

class Migration(migrations.Migration):
    dependencies = [
        ('trend', '06_ins_word_trends'),
    ]

    operations = [
        # 为 ProvincePatent 添加产业字段
        migrations.AddField(
            model_name='provincepatent',
            name='industry',
            field=models.CharField(
                max_length=100,
                verbose_name='产业',
                default='新材料',
                db_index=True
            ),
        ),
        
        # 为 Keyword 添加产业字段
        migrations.AddField(
            model_name='keyword',
            name='industry',
            field=models.CharField(
                max_length=100,
                verbose_name='产业',
                default='新材料',
                db_index=True
            ),
        ),
        
        # 为 TopOrg 添加产业字段
        migrations.AddField(
            model_name='toporg',
            name='industry',
            field=models.CharField(
                max_length=100,
                verbose_name='产业',
                default='新材料',
                db_index=True
            ),
        ),
        
        # 为 WordTrend 添加产业字段
        migrations.AddField(
            model_name='wordtrend',
            name='industry',
            field=models.CharField(
                max_length=100,
                verbose_name='产业',
                default='新材料',
                db_index=True
            ),
        ),
        
        # 添加复合索引
        migrations.AddIndex(
            model_name='provincepatent',
            index=models.Index(fields=['industry', 'province'], name='province_industry_idx'),
        ),
        migrations.AddIndex(
            model_name='keyword',
            index=models.Index(fields=['industry', 'count'], name='keyword_industry_idx'),
        ),
        migrations.AddIndex(
            model_name='toporg',
            index=models.Index(fields=['industry', 'count'], name='org_industry_idx'),
        ),
        migrations.AddIndex(
            model_name='wordtrend',
            index=models.Index(fields=['industry', 'year'], name='trend_industry_idx'),
        ),
        
        # 运行数据更新
        migrations.RunPython(
            update_industry_to_new_material,
            reverse_code=migrations.RunPython.noop
        ),
    ] 